#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <stdint.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

/**
 * \var g_run
 * \brief Running state of the program.
 */
static volatile sig_atomic_t g_run = 0;

/**
 * \brief Signal management.
 * \param code signal code
 */
static void signal_handler(int code)
{
  switch(code)
  {
    case SIGUSR1:
    case SIGUSR2:
    case SIGPIPE:
      break;
    case SIGINT:
    case SIGTERM:
      /* stop the program */
      g_run = 0;
      break;
    default:
      break;
  }
}

int socket_create(enum protocol_type type, const char* addr, uint16_t port,
    int reuse, int nodelay)
{
  int sock = -1;
  struct addrinfo hints;
  struct addrinfo* res = NULL;
  struct addrinfo* p = NULL;
  char service[8];

  snprintf(service, sizeof(service), "%u", port);
  service[sizeof(service)-1] = 0x00;

  memset(&hints, 0, sizeof(struct addrinfo));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = (type == TCP ? SOCK_STREAM : SOCK_DGRAM);
  hints.ai_protocol = (type == TCP ? IPPROTO_TCP : IPPROTO_UDP);
  hints.ai_flags = AI_PASSIVE;

  if(getaddrinfo(addr, service, &hints, &res) != 0)
  {
    return -1;
  }

  for(p = res ; p ; p = p->ai_next)
  {
    int on = 1;

    sock = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
    if(sock == -1)
    {
      continue;
    }

    if(reuse)
    {
      setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(int));
    }

    if (type == TCP && nodelay)
    {
      setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(int));
    }

    /* accept IPv6 and IPv4 on the same socket */
    on = 0;
    setsockopt(sock, IPPROTO_IPV6, IPV6_V6ONLY, &on, sizeof(int));

    if(bind(sock, p->ai_addr, p->ai_addrlen) == -1)
    {
      close(sock);
      sock = -1;
      continue;
    }

    /* socket bound, break the loop */
    break;
  }

  freeaddrinfo(res);
  p = NULL;

  return sock;
}

/**
 * \brief Entry point of the program.
 * \param argc number of argument
 * \param argv array of arguments
 * \return EXIT_SUCCESS or EXIT_FAILURE
 */
int main(int argc, char** argv)
{
  int sock = -1;
  struct sockaddr_storage addr;
  socklen_t addr_size = sizeof(struct sockaddr_storage);
  char buf[2500];
  ssize_t nb = -1;
  uint16_t port = 0;

  (void)argc; /* avoid compilation warning */

  signal(SIGUSR1, signal_handler);
  signal(SIGUSR2, signal_handler);
  signal(SIGPIPE, signal_handler);
  signal(SIGINT, signal_handler);
  signal(SIGTERM, signal_handler);

  port = argv[1] ? atol(argv[1]) : 4588;

  /* incorrect value */
  if(port == 0)
  {
    port = 4588;
  }

  /* try to bind on all addresses (IPv6+IPv4 mode) */
  sock = socket_create(IPPROTO_UDP, "::", port, 0, 0);

  if(sock == -1)
  {
    perror("socket");
    fprintf(stderr, "Maybe IPv6 is not available, try IPv4 only mode\n");
    sock = socket_create(IPPROTO_UDP, "0.0.0.0", port, 0, 1);
  }

  if(sock == -1)
  {
    perror("socket");
    exit(EXIT_FAILURE);
  }

  memset(&addr, 0x00, sizeof(struct sockaddr));
  memset(buf, 0x00, sizeof(buf));

  g_run = 1;

  fprintf(stdout, "UDP Echo server started on port %u\n", port);

  while(g_run)
  {
    fd_set fdsr;
    int nsock = sock;

    FD_ZERO(&fdsr);
    FD_SET(sock, &fdsr);

    nsock++;

    if(select(nsock, &fdsr, NULL, NULL, NULL) > 0)
    {
      if(FD_ISSET(sock, &fdsr))
      {
        nb = recvfrom(sock, buf, sizeof(buf), 0, (struct sockaddr*)&addr,
            &addr_size);
      }
    }
    else
    {
      perror("select");
      continue;
    }

    if(nb)
    {
      /* echo data received */
      if(sendto(sock, buf, nb, 0, (struct sockaddr*)&addr, addr_size) == -1)
      {
        perror("sendto");
      }
    }
    else
    {
      perror("recvfrom");
    }
  }

  close(sock);

  fprintf(stdout, "Exiting\n");

  return EXIT_SUCCESS;
}

